// Copyright (c) 2024 Huawei Technologies Co., Ltd.
// openFuyao is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//          http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

package lm

import (
	"bytes"
	"fmt"
	"os/exec"
	"strings"

	corev1 "k8s.io/api/core/v1"

	"openfuyao.com/npu-feature-discovery/internal/resources"
)

const (
	acceleratorTypeLabelKey = "accelerator-type"
)

func newAcceleratorTypeLabeler(node *corev1.Node) (Labeler, error) {
	cardNums, npuID, chipID, err := getCardNumsAndNPUInfo()
	if err != nil {
		return nil, fmt.Errorf("failed to get NPU info for node %s: %v", node.Name, err)
	}

	boardID, err := getBoardID(npuID, chipID)
	if err != nil {
		return nil, fmt.Errorf("failed to get board ID for node %s: %v", node.Name, err)
	}

	acceleratorType := getAcceleratorType(boardID, cardNums)
	if acceleratorType != "" {
		return Labels{acceleratorTypeLabelKey: acceleratorType}, nil
	}

	return nil, fmt.Errorf("no accelerator type found for node %s", node.Name)
}

func getAcceleratorType(boardID string, cardNums int) string {
	var accleratorType string

	switch {
	case contains(resources.CommonInfo["Atlas_800"], boardID):
		if cardNums == 8 {
			accleratorType = "module"
		} else if cardNums == 4 {
			accleratorType = "half"
		}
	case contains(resources.CommonInfo["Atlas_800_A2"], boardID),
		contains(resources.CommonInfo["Atlas_900_A2_PoD"], boardID):
		accleratorType = "module-910b-8"
	case contains(resources.CommonInfo["Atlas_200T_A2_Box16"], boardID):
		accleratorType = "module-910b-16"
	case contains(resources.CommonInfo["Atlas_300T"], boardID):
		accleratorType = "card"
	case contains(resources.CommonInfo["Atlas_300T_A2"], boardID):
		accleratorType = "card-910b-2"
	}

	return accleratorType
}

func contains(slice []string, item string) bool {
	for _, s := range slice {
		if s == item {
			return true
		}
	}
	return false
}

func getBoardID(npuID, chipID string) (string, error) {
	cmdOutput, err := executeCommand(fmt.Sprintf("/usr/local/sbin/npu-smi info -t board -i %s -c %s", npuID, chipID))
	if err != nil {
		return "", err
	}

	var boardID string
	for _, line := range strings.Split(cmdOutput, "\n") {
		if strings.Contains(line, "Board") && strings.Contains(line, ":") {
			line = strings.TrimSpace(line)
			boardID = strings.Split(line, ":")[1]
			boardID = strings.TrimSpace(boardID)
			break
		}
	}

	return strings.ToLower(boardID), nil
}

func getCardNumsAndNPUInfo() (int, string, string, error) {
	cmdOutput, err := executeCommand("/usr/local/sbin/npu-smi info -m")
	if err != nil {
		return 0, "", "", err
	}

	cardNums := 0
	npuID, chipID := "0", "0"

	for _, line := range strings.Split(cmdOutput, "\n") {
		line = strings.ReplaceAll(line, "\t", "")
		if strings.Contains(line, "Ascend") {
			parts := strings.Fields(line)
			if len(parts) >= 3 {
				cardNums++
				if cardNums == 1 {
					npuID, chipID = parts[0], parts[1]
				}
			}
		}
	}
	return cardNums, npuID, chipID, nil
}

func executeCommand(cmd string) (string, error) {
	cmdArgs := []string{
		"-i/proc/1/ns/ipc",
		"-m/proc/1/ns/mnt",
		"-n/proc/1/ns/net",
	}

	cmdParts := strings.Fields(cmd)

	cmdArgs = append(cmdArgs, cmdParts...)

	cmdOut := exec.Command("nsenter", cmdArgs...)
	var out bytes.Buffer
	cmdOut.Stdout = &out
	cmdOut.Stderr = &out

	err := cmdOut.Run()
	if err != nil {
		return "", fmt.Errorf("error executing command: %w\nOutput: %s", err, out.String())
	}

	return out.String(), nil
}
